import gradio as gr
import torch
import torch.nn as nn
from torchvision import transforms
from PIL import Image

class BasicBlock(nn.Module):
    expansion = 1
 
    def __init__(self, in_channel, out_channel, stride=1, downsample=None):
        super(BasicBlock, self).__init__()
 
        self.conv1 = nn.Conv2d(in_channel, out_channel, kernel_size=3, padding=1, stride=stride, bias=False)
        self.bn1 = nn.BatchNorm2d(out_channel)    
        self.conv2 = nn.Conv2d(out_channel, out_channel, kernel_size=3, padding=1, bias=False)
        self.bn2 = nn.BatchNorm2d(out_channel)

        self.downsample = downsample
        self.relu = nn.ReLU(inplace=True)
 
    def forward(self, X):
        identity = X
        Y = self.relu(self.bn1(self.conv1(X)))
        Y = self.bn2(self.conv2(Y))
 
        if self.downsample is not None:    
            identity = self.downsample(X)
 
        return self.relu(Y + identity)
 
 
class BottleNeck(nn.Module):
    expansion = 4

    def __init__(self, in_channel, out_channel, stride=1, downsample=None):
        super(BottleNeck, self).__init__()

        self.conv1 = nn.Conv2d(in_channel, out_channel, kernel_size=1, bias=False)
        self.bn1 = nn.BatchNorm2d(out_channel)    
        self.conv2 = nn.Conv2d(out_channel, out_channel, kernel_size=3, stride=stride, padding=1, bias=False)
        self.bn2 = nn.BatchNorm2d(out_channel)
        self.conv3 = nn.Conv2d(out_channel, out_channel * self.expansion, kernel_size=1, bias=False)
        self.bn3 = nn.BatchNorm2d(out_channel * self.expansion)  
 
        self.downsample = downsample
        self.relu = nn.ReLU(inplace=True)
 
    def forward(self, X):
        identity = X
 
        Y = self.relu(self.bn1(self.conv1(X)))
        Y = self.relu(self.bn2(self.conv2(Y)))
        Y = self.bn3(self.conv3(Y))
 
        if self.downsample is not None:    
            identity = self.downsample(X)
 
        return self.relu(Y + identity)
 
 
class ResNet(nn.Module):

    def __init__(self, residual, num_residuals, num_classes=10, include_top=True):
        super(ResNet, self).__init__()
 
        self.out_channel = 64    
        self.include_top = include_top
 
        self.conv1 = nn.Conv2d(3, self.out_channel, kernel_size=7, stride=2, padding=3,
                               bias=False)    
        self.bn1 = nn.BatchNorm2d(self.out_channel)
        self.relu = nn.ReLU(inplace=True)
        self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1)
        self.conv2 = self.residual_block(residual, 64, num_residuals[0])
        self.conv3 = self.residual_block(residual, 128, num_residuals[1], stride=2)
        self.conv4 = self.residual_block(residual, 256, num_residuals[2], stride=2)
        self.conv5 = self.residual_block(residual, 512, num_residuals[3], stride=2)
        if self.include_top:
            self.avgpool = nn.AdaptiveAvgPool2d((1, 1))
            self.fc = nn.Linear(512 * residual.expansion, num_classes)
 
        # 对conv层进行初始化操作
        for m in self.modules():
            if isinstance(m, nn.Conv2d):
                nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')
            elif isinstance(m, (nn.BatchNorm2d, nn.GroupNorm)):
                nn.init.constant_(m.weight, 1)
                nn.init.constant_(m.bias, 0)
 
    def residual_block(self, residual, channel, num_residuals, stride=1):
        downsample = None
 
        if stride != 1 or self.out_channel != channel * residual.expansion:
            downsample = nn.Sequential(
                nn.Conv2d(self.out_channel, channel * residual.expansion, kernel_size=1, stride=stride, bias=False),
                nn.BatchNorm2d(channel * residual.expansion))
 
        block = []
        block.append(residual(self.out_channel, channel, downsample=downsample, stride=stride))
        self.out_channel = channel * residual.expansion
 
        for _ in range(1, num_residuals):
            block.append(residual(self.out_channel, channel))
 
        return nn.Sequential(*block)
 

    def forward(self, X):
        Y = self.relu(self.bn1(self.conv1(X)))
        Y = self.maxpool(Y)
        Y = self.conv5(self.conv4(self.conv3(self.conv2(Y))))
 
        if self.include_top:
            Y = self.avgpool(Y)
            Y = torch.flatten(Y, 1)
            Y = self.fc(Y)
 
        return Y
 
 
 
# 构建ResNet模型
def resnet18(num_classes=10, include_top=True):
    return ResNet(BottleNeck, [2, 2, 2, 2], num_classes=num_classes, include_top=include_top)
def resnet50(num_classes=10, include_top=True):
    return ResNet(BottleNeck, [3, 4, 6, 3], num_classes=num_classes, include_top=include_top)
def resnet101(num_classes=10, include_top=True):
    return ResNet(BottleNeck, [3, 4, 23, 3], num_classes=num_classes, include_top=include_top)
def resnet152(num_classes=10, include_top=True):
    return ResNet(BottleNeck, [3, 8, 36, 3], num_classes=num_classes, include_top=include_top)

# 加载训练好的模型
model = resnet101(num_classes=10)  # 修改类别数
model.load_state_dict(torch.load('./models/resnet101_best_model.pth'))
model.eval()

# 定义预处理转换
transform = transforms.Compose([
    transforms.Resize(256),
    transforms.CenterCrop(224),
    transforms.ToTensor(),
    transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
])

# 定义类别标签
class_names = ['Bike', 'Bird', 'Car', 'Person', 'Plane', 
               'Ship', 'Traffic light', 'Train', 'Truck', 'Zebra crossing']
resize_image = 224

def predict_images(file_paths):
    processed_images = []
    labels = []
    
    for file_path in file_paths:
        # 打开图像
        img = Image.open(file_path)
        
        # 保存原始图像用于显示
        img_resized = img.resize((resize_image, resize_image), Image.Resampling.LANCZOS)
        
        # 预处理并预测
        img_tensor = transform(img_resized).unsqueeze(0)
        with torch.no_grad():
            outputs = model(img_tensor)
        _, preds = torch.max(outputs, 1)
        
        processed_images.append(img_resized)
        labels.append(class_names[preds[0]])
    labels_text = ",\n".join(labels)
    
    return processed_images, labels_text
 
# 创建Gradio界面
with gr.Blocks() as demo:
    gr.Markdown("# 图像分类器")
    with gr.Row():
        input_imgs = gr.Files(label="输入图像", file_types=["image"])  # 指定只接受图像文件
        output_imgs = gr.Gallery(label="处理后图像")
        labels = gr.Label(label="预测类别")
    with gr.Row():
        submit_btn = gr.Button("识别")
    submit_btn.click(
        fn=predict_images,
        inputs=input_imgs,
        outputs=[output_imgs, labels]
    )

if __name__ == "__main__":
    demo.launch()